﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using SecureConnection.Encryption;
using SecureConnection.Optimization;
using SecureConnection.Optimization.Compression;
using SecureConnection.Packets;
using SecureConnection.Utils;

namespace SecureConnection
{
    public class Connection
    {
        #region Private/Internal Variables
        private static List<int> _allocatedIds = new List<int>();

        private Socket _socket;

        internal object _socketOperation = new object();
        internal object _receiveFlow = new object();
        internal object _sendFlow = new object();

        private byte[] _encryptionKey;
        private bool _doKeyExchange;

        internal DynamicKeyEncryption _dynamicEnc;
        internal QuickLZ _compression;

        private bool _isReceivingPayload = false;
        private BasicBuffer _headerBuffer;
        private BasicBuffer _payloadBuffer;
        private bool _sendError = false;

        private byte _receiveIndex = 0;
        private byte _sendIndex = 0;

        //private bool headerParsed = false;
        //private int extraLength = 0;
        //private BasicBuffer receiveBuffer;
        //private BasicBuffer packetBuffer;
        #endregion

        #region Public Variables
        public const int HeaderSize = 5;

        public bool Listening { get; private set; }
        public bool Connected { get { return _socket != null ? _socket.Connected : false; } }
        public bool IsNull { get { return _socket == null; } }

        public ulong BytesSent { get; private set; }
        public ulong BytesReceived { get; private set; }

        public int SharedId {get; internal set;}


        public EndPoint RemoteEndPoint { get { return Connected ? _socket.RemoteEndPoint : null; } }
        public EndPoint LocalEndPoint { get { return Connected ? _socket.LocalEndPoint : null; } }
        #endregion

        #region Constructors
        private void InitializeVariables(byte[] key, bool dokeyExchange)
        {
            this.Listening = false;
            this.SharedId = 0;
            this.BytesReceived = 0;
            this.BytesSent = 0;

            if(!dokeyExchange && (key == null || key.Length == 0)) throw new Exception("Empty key!");

            _encryptionKey = key;
            _doKeyExchange = dokeyExchange;
            _compression = new QuickLZ();

            _socket.NoDelay = true;
        }

        public Connection(bool doKeyExchange = true, byte[] key = null)
        {
            _socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);

            InitializeVariables(key, doKeyExchange);
        }
        private Connection(Socket s, bool doKeyExchange, byte[] key)
        {
            _socket = s;

            InitializeVariables(key, doKeyExchange);
        }
        #endregion

        #region Handshake

        //Thanks to Grimlock, http://www.codeproject.com/Tips/257269/Find-Prime-Numbers-in-Csharp-Quickly
        private static int[] GeneratePrimeTable(int bound)
        {
            if (bound < 2) throw new Exception();
            List<int> primes = new List<int>();
            primes.Add(2);
            BitArray notPrime = new BitArray((bound - 1) >> 1);
            int limit = ((int)(Math.Sqrt(bound)) - 1) >> 1;
            for (int i = 0; i < limit; i++)
            {
                if (notPrime[i]) continue;
                int prime = 2 * i + 3;
                primes.Add(prime);
                for (int j = (prime * prime - 2) >> 1; j < notPrime.Count; j += prime)
                {
                    notPrime[j] = true;
                }
            }
            for (int i = limit; i < notPrime.Count; i++)
            {
                if (!notPrime[i]) primes.Add(2 * i + 3);
            }

            return primes.ToArray();
        }

        private bool SendHandshake()
        {
            try
            {
                lock (_socketOperation)
                {
                    Random rnd = new Random(new Guid().GetHashCode());
                    byte[] buffer = null;

                    _socket.Send(new byte[] { _doKeyExchange ? (byte)1 : (byte) 0});
                    if (_doKeyExchange)
                    {
                        //Diffie Hellman key exchange algorithm: http://en.wikipedia.org/wiki/Diffie%E2%80%93Hellman_key_exchange

                        byte[] key = new byte[512]; //4096 bit
                        int offset = 0;

                        buffer = new byte[4];
                        int[] primes = GeneratePrimeTable(1000000);

                        for (int i = 0; i < key.Length / 4; i++)
                        {
                            int p = primes[rnd.Next(0, primes.Length)];
                            int g = rnd.Next(1, p - 1);
                            int a = rnd.Next(1, p - 1);
                            _socket.Send(BitConverter.GetBytes(p));
                            _socket.Send(BitConverter.GetBytes(g));
                            int A = (int)(g ^ a % p);
                            _socket.Send(BitConverter.GetBytes(A));
                            if (_socket.Receive(buffer) != 4)
                                return false;
                            int B = BitConverter.ToInt32(buffer, 0);

                            int s = (int)(B ^ a % p);

                            Buffer.BlockCopy(BitConverter.GetBytes(s), 0, key, offset, 4);
                            offset += 4;
                        }


                        _encryptionKey = key;
                    }
                    _dynamicEnc = new DynamicKeyEncryption(_encryptionKey);

                    byte[] randomBuffer = new byte[1024];

                    rnd.NextBytes(randomBuffer);
                    _socket.Send(randomBuffer);

                    buffer = new byte[1024];
                    if (_socket.Receive(buffer) != 1024)
                    { Close(CloseReason.VerificationFailed); return false; }
                    buffer = _dynamicEnc.Decrypt(buffer);

                    for (int i = 0; i < 1024; i++)
                        if (buffer[i] != randomBuffer[i])
                        { Close(CloseReason.VerificationFailed); return false; }

                    int counter = 3;
                    buffer = new byte[1];
                    while (true)
                    {
                        lock (_allocatedIds)
                        {
                            this.SharedId = rnd.Next(int.MinValue, int.MaxValue);
                            while (this.SharedId == 0 || _allocatedIds.Contains(SharedId))
                                this.SharedId = rnd.Next(int.MinValue, int.MaxValue);
                            _allocatedIds.Add(this.SharedId);
                        }

                        _socket.Send(_dynamicEnc.Encrypt(BitConverter.GetBytes(this.SharedId)));

                        if (_socket.Receive(buffer) != 1)
                        { Close(CloseReason.VerificationFailed); return false; }

                        if (buffer[0] == 1)
                            break;

                        counter--;
                        if (counter == 0)
                        { Close(CloseReason.VerificationFailed); return false; }
                    }
                    return true;
                }
            }
            catch (Exception) { Close(CloseReason.VerificationFailed); return false; }
        }

        private static bool AcceptHandshake(Connection conn)
        {
            try
            {
                lock (conn._socketOperation)
                {
                    Random rnd = new Random(DateTime.Now.Millisecond);
                    byte[] buffer;

                    buffer = new byte[1];
                    if (conn._socket.Receive(buffer) != 1)
                    { conn.Close(CloseReason.VerificationFailed); return false; }

                    if ((buffer[0] == 1 ? true : false) != conn._doKeyExchange)
                    { conn.Close(CloseReason.VerificationFailed); return false; }

                    if (conn._doKeyExchange)
                    {
                        //Diffie Hellman key exchange algorithm: http://en.wikipedia.org/wiki/Diffie%E2%80%93Hellman_key_exchange

                        buffer = new byte[4];

                        byte[] key = new byte[512]; //4096 bit
                        int offset = 0;
                        for (int i = 0; i < key.Length / 4; i++)
                        {
                            if (conn._socket.Receive(buffer) != 4)
                            { conn.Close(CloseReason.VerificationFailed); return false; }
                            int p = BitConverter.ToInt32(buffer, 0);
                            if (conn._socket.Receive(buffer) != 4)
                            { conn.Close(CloseReason.VerificationFailed); return false; }
                            int g = BitConverter.ToInt32(buffer, 0);
                            if (conn._socket.Receive(buffer) != 4)
                            { conn.Close(CloseReason.VerificationFailed); return false; }
                            int A = BitConverter.ToInt32(buffer, 0);

                            int b = rnd.Next(1, p - 1);
                            int B = g ^ b % p;
                            conn._socket.Send(BitConverter.GetBytes(B));

                            int s = A ^ b % p;

                            Buffer.BlockCopy(BitConverter.GetBytes(s), 0, key, offset, 4);
                            offset += 4;
                        }
                        conn._encryptionKey = key;
                    }
                    conn._dynamicEnc = new DynamicKeyEncryption(conn._encryptionKey);

                    buffer = new byte[1024];
                    if (conn._socket.Receive(buffer) != 1024)
                    { conn.Close(CloseReason.VerificationFailed); return false; }

                    buffer = conn._dynamicEnc.Encrypt(buffer);
                    conn._socket.Send(buffer);

                    int counter = 3;
                    buffer = new byte[4];
                    while(true)
                    {
                        if (conn._socket.Receive(buffer) != 4)
                        { conn.Close(CloseReason.VerificationFailed); return false; }

                        conn.SharedId = BitConverter.ToInt32(conn._dynamicEnc.Decrypt(buffer), 0);

                        lock (_allocatedIds)
                        {
                            if (!_allocatedIds.Contains(conn.SharedId))
                                break;
                        }
                        conn._socket.Send(new byte[] { 0 });

                        counter--;
                        if (counter == 0)
                        { conn.Close(CloseReason.VerificationFailed); return false; }
                    }
                    conn._socket.Send(new byte[] { 1 });

                    return true;
                }
            }
            catch (Exception) { conn.Close(CloseReason.VerificationFailed); return false; }
        }
        #endregion

        #region Connecting
        public void Connect(string hostIp, ushort hostPort)
        {
            lock (_socketOperation)
            {
                if (IsNull)
                    throw new Exception("Socket already used");
                if (Listening)
                    throw new Exception("This is a listening socket");
                if (Connected)
                    throw new Exception("Already connected");

                _socket.Connect(hostIp, (int)hostPort);

                if (!SendHandshake())
                    throw new Exception("Verification failed");

                BeginReceive();
            }
        }
        #endregion    

        #region Disconnecting
        public void Close()
        {
            Close(CloseReason.User);
        }

        private void Close(CloseReason reason)
        {
            lock (_socketOperation)
            {
                if (!IsNull)
                {
                        try
                        {
                            if (_socket != null)
                                _socket.Shutdown(SocketShutdown.Both);
                        }
                        catch { }
                        try
                        {
                            if (_socket != null)
                                _socket.Close();
                        }
                        catch { }
                        _socket = null;

                        Listening = false;

                        OnDisconnect(reason, this);
                }
            }
        }
        #endregion

        #region Listening and Accepting
        public void Listen(ushort localPort, int backLog)
        {
            lock (_socketOperation)
            {
                if (IsNull)
                    throw new Exception("Socket already used");
                if (Listening)
                    throw new Exception("Already listening");
                if (Connected)
                    throw new Exception("Socket connected");

                _socket.Bind(new IPEndPoint(IPAddress.Any, localPort));
                _socket.Listen(backLog);

                Listening = true;
            }
        }

        public Connection Accept()
        {
            lock (_socketOperation)
            {
                if (IsNull)
                    throw new Exception("Socket already used");
                if (!Listening)
                    throw new Exception("Socket not listening");
                if (Connected)
                    throw new Exception("Socket connected");

                Connection conn = new Connection(_socket.Accept(), _doKeyExchange, _encryptionKey);

                if (!AcceptHandshake(conn))
                    throw new Exception("Verification failed");

                conn.BeginReceive();
                return conn;
            }
        }

        public void BeginAccept(Callback<Connection> callBack )
        {
            lock (_socketOperation)
            {
                if (IsNull)
                    throw new Exception("Socket already used");
                if (!Listening)
                    throw new Exception("Socket not listening");
                if (Connected)
                    throw new Exception("Socket connected");

                _socket.BeginAccept(AcceptCallback, callBack);
            }
        }
        private void AcceptCallback(IAsyncResult ar)
        {
            lock (_socketOperation)
            {
                Callback<Connection> callback = (Callback<Connection>)ar.AsyncState;
                try
                {
                    Connection conn = new Connection(_socket.EndAccept(ar), _doKeyExchange, _encryptionKey);

                    if (!AcceptHandshake(conn))
                    {
                        BeginAccept(callback);
                        return;
                    }

                    conn.BeginReceive();
                    callback(conn, false, null, this);
                }
                catch (Exception ex) { callback(null, true, ex, this); }
            }
        }
        #endregion

        #region Receiving

        //private void BeginReceive()
        //{
        //    lock (SocketOperation)
        //    {
        //        packetBuffer = new BasicBuffer();
        //        receiveBuffer = new BasicBuffer();

        //        receiveBuffer.Resize(BufferSize);
        //        tcpSocket.BeginReceive(receiveBuffer.Payload, receiveBuffer.Offset, receiveBuffer.Payload.Length, SocketFlags.None, TcpReceiveCallback, null);
        //    }
        //}

        private void BeginReceive()
        {
            lock (_socketOperation)
            {
                _headerBuffer = new BasicBuffer();
                _payloadBuffer = new BasicBuffer();

                _headerBuffer.Resize(HeaderSize);
                _socket.BeginReceive(_headerBuffer.Payload, _headerBuffer.Offset, _headerBuffer.Size, SocketFlags.None, TcpReceiveCallback, null);
            }
        }

        //private unsafe void TcpReceiveCallback(IAsyncResult ar)
        //{
        //    lock (ReceiveFlow)
        //    {
        //        int br = tcpSocket.EndReceive(ar);
        //        receiveBuffer.Size = br + extraLength;
        //        extraLength = 0;

        //        if (br == 0)
        //        { Close(); return; }

        //        while (!receiveBuffer.EndOfBuffer)
        //        {
        //            if (!headerParsed)
        //            {
        //                if (receiveBuffer.Left < HeaderSize)
        //                {
        //                    int copyLength = receiveBuffer.Left;

        //                    Buffer.BlockCopy(receiveBuffer.Payload, receiveBuffer.Offset, receiveBuffer.Payload, 0, copyLength);

        //                    extraLength = copyLength;
        //                    tcpSocket.BeginReceive(receiveBuffer.Payload, extraLength, receiveBuffer.Payload.Length - extraLength, SocketFlags.None, TcpReceiveCallback, null);
        //                }
        //                else
        //                {
        //                    dynamicEnc.Decrypt(ref receiveBuffer.Payload, receiveBuffer.Offset, HeaderSize);

        //                    int payloadSize = BitConverter.ToInt32(receiveBuffer.Payload, receiveBuffer.Offset);

        //                    byte sendIndex = receiveBuffer.Payload[receiveBuffer.Offset + 4];
        //                    if (sendIndex != receiveIndex)
        //                    { Close(); return; }
        //                    receiveIndex = (byte)((receiveIndex + 1) % 256);

        //                    if (payloadSize < 1 || payloadSize > 5242880)
        //                    { Close(); return; }

        //                    packetBuffer.Resize(payloadSize);
        //                    headerParsed = true;

        //                    receiveBuffer.Offset += HeaderSize;
        //                }
        //            }
        //            else
        //            {
        //                int copyLength = receiveBuffer.Left > packetBuffer.Left ? packetBuffer.Left : receiveBuffer.Left;

        //                Buffer.BlockCopy(receiveBuffer.Payload, receiveBuffer.Offset, packetBuffer.Payload, packetBuffer.Offset, copyLength);

        //                packetBuffer.Offset += copyLength;
        //                receiveBuffer.Offset += copyLength;

        //                if (packetBuffer.EndOfBuffer)
        //                {
        //                    dynamicEnc.Decrypt(ref packetBuffer.Payload, 0, packetBuffer.Size);
        //                   // packetBuffer.Payload = compression.Decompress(packetBuffer.Payload);

        //                    OnData(packetBuffer.Payload, this);

        //                    headerParsed = false;
        //                }
        //            }
        //        }

        //        receiveBuffer.Reset();
        //        tcpSocket.BeginReceive(receiveBuffer.Payload, receiveBuffer.Offset, receiveBuffer.Payload.Length, SocketFlags.None, TcpReceiveCallback, null);
        //    }
        //}

        private unsafe void TcpReceiveCallback(IAsyncResult ar)
        {
            lock (_receiveFlow)
            {
                try
                {
                    int br = _socket.EndReceive(ar);

                    if (br == 0)
                    { Close(CloseReason.BytesReadZero); return; }

                    BytesReceived += (ulong)br;

                    if (!_isReceivingPayload)
                    {
                        _headerBuffer.Offset += br;

                        if (br < HeaderSize)
                            _socket.BeginReceive(_headerBuffer.Payload, _headerBuffer.Offset, _headerBuffer.Size - _headerBuffer.Offset, SocketFlags.None, TcpReceiveCallback, null);
                        else if (br > HeaderSize)
                        {
                            Close(CloseReason.IncorrectBytesRead); return;
                        }

                        _dynamicEnc.Decrypt(ref _headerBuffer.Payload, 0, _headerBuffer.Size);

                        int payloadSize = BitConverter.ToInt32(_headerBuffer.Payload, 0);

                        if (payloadSize < 1 || payloadSize > 5242880)
                        { Close(CloseReason.HeaderInvalid); return; }

                        byte sendIndex = _headerBuffer.Payload[4];

                        if (sendIndex != _receiveIndex)
                        { Close(CloseReason.HeaderInvalid); return; }

                        _receiveIndex = (byte)((_receiveIndex + 1) % 256);

                        _payloadBuffer.Resize(payloadSize);
                        _isReceivingPayload = true;

                        _socket.BeginReceive(_payloadBuffer.Payload, _payloadBuffer.Offset, _payloadBuffer.Size, SocketFlags.None, TcpReceiveCallback, null);
                    }
                    else
                    {
                        _payloadBuffer.Offset += br;

                        if (_payloadBuffer.EndOfBuffer)
                        {
                            _dynamicEnc.Decrypt(ref _payloadBuffer.Payload, 0, _payloadBuffer.Size, HeaderSize);
                            _payloadBuffer.Payload = _compression.Decompress(_payloadBuffer.Payload);

                            OnData(_payloadBuffer.Payload, this);

                            _isReceivingPayload = false;

                            _headerBuffer.Reset();
                            _payloadBuffer.Reset();
                            _socket.BeginReceive(_headerBuffer.Payload, _headerBuffer.Offset, _headerBuffer.Size, SocketFlags.None, TcpReceiveCallback, null);
                        }
                        else _socket.BeginReceive(_payloadBuffer.Payload, _payloadBuffer.Offset, _payloadBuffer.Size - _payloadBuffer.Offset, SocketFlags.None, TcpReceiveCallback, null);
                    }
                }
                catch (Exception) { Close(CloseReason.Unkown); return; }
            }
        }
        #endregion

        #region Sending
        public void Send(byte[] data)
        {
            try
            {
                lock (_sendFlow)
                {
                    if (IsNull)
                        throw new Exception("Socket already used");
                    if (Listening)
                        throw new Exception("Socket is listening");
                    if (!Connected)
                        throw new Exception("Socket not connected");

                    data = _compression.Compress(data, 1);

                    using (MemoryStream ms = new MemoryStream())
                    {
                        ms.Write(BitConverter.GetBytes(data.Length), 0, 4);
                        ms.WriteByte(_sendIndex);
                        ms.Write(data, 0, data.Length);


                        if (!_socket.BeginSend(_dynamicEnc.Encrypt(ms.ToArray()), 0, (int)ms.Length, SocketFlags.None, SendCallback, null)
                          .AsyncWaitHandle.WaitOne(15000))
                        { Close(CloseReason.SendTimeout); throw new Exception("Timed out"); }

                        if (_sendError)
                        { Close(CloseReason.SendError); throw new Exception("Error"); }

                        _sendIndex = (byte)((_sendIndex + 1) % 256);

                        BytesSent += (ulong)ms.Length;
                    }

                }
            }
            catch (Exception)
            {
                Close(CloseReason.Unkown);
                throw;
            }
        }
        private void SendCallback(IAsyncResult ar)
        {
            try
            {
                _socket.EndSend(ar);
            }
            catch { _sendError = true; }
        }

        #endregion

        #region API's, Delegates and Events
        private delegate void SendDelegate(byte[] data);

        public delegate void Callback<T>(T ret, bool err, Exception ex, Connection sender);
        public delegate void Callback(bool err, Exception ex, Connection sender);

        public delegate void OnDataHandler(byte[] data, Connection sender);
        public event OnDataHandler OnData = null;

        public delegate void OnDisconnectHandler(CloseReason reason, Connection sender);
        public event OnDisconnectHandler OnDisconnect = delegate { };
        #endregion
    }
    public enum CloseReason
    {
        User,
        ReceiveException,
        BytesReadZero,
        PacketVerificationFailed,
        IncorrectBytesRead,
        HeaderInvalid,
        SendTimeout,
        SendError,
        VerificationFailed,
        Unkown
    }
}
